import tkinter as tk
from tkinter import *
from tkinter import filedialog
import tkinter.messagebox
from PIL import Image, ImageTk
from PIL import ImageDraw
from PIL import ImageFont
import re
import os

'''松耦合'''
# 加密弹窗
class MyDialogEncry(tk.Toplevel):
  def __init__(self):
    super().__init__()
    self.title('图像加密')
    self.setup_UI_encry()

  def setup_UI_encry(self):
    global filePath1, filePath2

    # 生成画布
    frm = tk.Frame(self)
    frm.pack(padx=8, pady=8, ipadx=4, ipady=0)
    # 画布布局
    tk.Label(frm, text='原始图像文件路径：').grid(row=0, column=0, padx=5, pady=5)
    filePath1 = StringVar()
    filepath_en1 = tk.Entry(frm, textvariable=filePath1)
    filepath_en1.grid(row=0, column=1, padx=5, pady=5)
    tk.Button(frm, text="添加图像", default='active', command=self.show_path1).grid(row=0, column=2, padx=5, pady=5)
    # 画布布局
    tk.Label(frm, text='水印图像文件路径：').grid(row=1, column=0, padx=5, pady=5)
    filePath2 = StringVar()
    filepath_en2 = tk.Entry(frm, textvariable=filePath2)
    filepath_en2.grid(row=1, column=1, padx=5, pady=5)
    tk.Button(frm, text="添加图像", default='active', command=self.show_path2).grid(row=1, column=2,padx=5, pady=5)
    # 生成画布
    frm2 = tk.Frame(self)
    frm2.pack(padx=8, pady=8, ipadx=4, ipady=0, fill = 'x')
    tk.Button(frm2, text="取消", default='active', command=self.cancel).pack(padx=5, pady=5, side=tk.RIGHT)
    tk.Button(frm2, text="确定", default='active', command=self.ok).pack(padx=5, pady=5, side=tk.RIGHT)

  # 获取待处理图片的文件路径
  def show_path1(self):
      global img1, filePath, file_NewName, file_dir, name
      # 获取文件的绝对路径
      filePath = filedialog.askopenfilename(defaultextension=".bmp", filetypes=[('image', '.bmp'), ("BMP", ".bmp")])
      file_dir = os.path.dirname(filePath)
      file_name = os.path.basename(filePath)
      # 删jpg后缀
      name = re.sub(re.escape(".bmp") + '$', '', file_name)

      # 获取打开的图片
      img1 = Image.open(filePath, 'r')
      filePath1.set(filePath)

  # 获取待处理图片的文件路径
  def show_path2(self):
      global img2
      # 获取文件的绝对路径
      filePath = filedialog.askopenfilename(defaultextension=".bmp", filetypes=[('image', '.bmp'), ("BMP", ".bmp")])
      img2 = Image.open(filePath, 'r')
      filePath2.set(filePath)

  def plus(self, str):
      # 返回指定长度的字符串，原字符串右对齐，前面填充0。
      return str.zfill(8)

  def getCode(self,img):
      str = ""
      # 获取到水印的宽和高进行遍历
      for i in range(img.size[0]):
          for j in range(img.size[1]):
              # 获取水印的每个像素值
              rgb = img.getpixel((i, j))
              # 将像素值转为二进制后保存
              str = str + self.plus(bin(rgb[0]).replace('0b', ''))
              str = str + self.plus(bin(rgb[1]).replace('0b', ''))
              str = str + self.plus(bin(rgb[2]).replace('0b', ''))
              # print(plus(bin(rgb[0]).replace('0b', ''))+"\n")
              # print(plus(bin(rgb[1]).replace('0b', '')) + "\n")
              # print(plus(bin(rgb[2]).replace('0b', '')) + "\n")
      # print(str)
      return str

  def encry(self, img, code):
      # 计数器
      count = 0
      # 二进制像素值的长度，可以认为要写入图像的文本长度，提取（解密）时也需要此变量
      codeLen = len(code)
      print(codeLen)

      # 获取到图像的宽、高进行遍历
      for i in range(img.size[0]):
          for j in range(img.size[1]):
              # 获取到图片的每个像素值
              data = img.getpixel((i, j))

              # 如果计数器等于长度，代表水印写入完成
              if count == codeLen:
                  break

              # 将获取到的RGB数值分别保存
              r = data[0]
              g = data[1]
              b = data[2]

              """
              下面的是像素值替换，通过取模2得到最后一位像素值（0或1），
              然后减去最后一位像素值，在将code的值添加过来
              """

              r = (r - r % 2) + int(code[count])
              count += 1
              if count == codeLen:
                  img.putpixel((i, j), (r, g, b))
                  break

              g = (g - g % 2) + int(code[count])
              count += 1
              if count == codeLen:
                  img.putpixel((i, j), (r, g, b))
                  break

              b = (b - b % 2) + int(code[count])
              count += 1
              if count == codeLen:
                  img.putpixel((i, j), (r, g, b))
                  break

              # 每3次循环表示一组RGB值被替换完毕，可以进行写入
              if count % 3 == 0:
                  img.putpixel((i, j), (r, g, b))

      # 要保存的文件名
      self.path = file_dir + '/' + name + '加密图像' + '.bmp'
      # 保存图片
      img.save(self.path)
      # 显示图片
      img.show()

  def ok(self):
    # 将图像转换为RGB通道，才能分别获取R,G,B的值
    rgb_im1 = img1.convert('RGB')
    rgb_im2 = img2.convert('RGB')
    # 将水印的像素值转为文本
    code = self.getCode(rgb_im2)
    # 将水印写入图像
    self.encry(rgb_im1, code)
    # 设置数据
    self.userinfo = [self.path]
    # 销毁窗口
    self.destroy()

  def cancel(self):
    self.userinfo = None # 空！
    self.destroy()


# 解密弹窗
class MyDialogDeEncry(tk.Toplevel):
    def __init__(self):
        super().__init__()
        self.title('水印检测')
        self.setup_UI_deEncry()

    def setup_UI_deEncry(self):
        global filePath1, filePath2

        # 生成画布
        frm = tk.Frame(self)
        frm.pack(padx=8, pady=8, ipadx=4, ipady=0)
        # 画布布局
        tk.Label(frm, text='加密图像文件路径：').grid(row=0, column=0, padx=5, pady=5)
        filePath1 = StringVar()
        filepath_en1 = tk.Entry(frm, textvariable=filePath1)
        filepath_en1.grid(row=0, column=1, padx=5, pady=5)
        tk.Button(frm, text="添加图像", default='active', command=self.show_path1).grid(row=0, column=2, padx=5, pady=5)
        # 画布布局
        tk.Label(frm, text='水印图像文件路径：').grid(row=1, column=0, padx=5, pady=5)
        filePath2 = StringVar()
        filepath_en2 = tk.Entry(frm, textvariable=filePath2)
        filepath_en2.grid(row=1, column=1, padx=5, pady=5)
        tk.Button(frm, text="添加图像", default='active', command=self.show_path2).grid(row=1, column=2, padx=5, pady=5)
        # 生成画布
        frm2 = tk.Frame(self)
        frm2.pack(padx=8, pady=8, ipadx=4, ipady=0, fill='x')
        tk.Button(frm2, text="取消", default='active', command=self.cancel).pack(padx=5, pady=5, side=tk.RIGHT)
        tk.Button(frm2, text="确定", default='active', command=self.ok).pack(padx=5, pady=5, side=tk.RIGHT)

    # 获取待处理图片的文件路径
    def show_path1(self):
        global img1, filePath, file_NewName, file_dir, name
        # 获取文件的绝对路径
        filePath = filedialog.askopenfilename(defaultextension=".bmp", filetypes=[('image', '.bmp'), ("BMP", ".bmp")])
        file_dir = os.path.dirname(filePath)
        file_name = os.path.basename(filePath)
        # 删jpg后缀
        name = re.sub(re.escape(".bmp") + '$', '', file_name)

        # 获取打开的图片
        img1 = Image.open(filePath, 'r')
        filePath1.set(filePath)

    # 获取待处理图片的文件路径
    def show_path2(self):
        global img2, image_width, image_height, image_length
        # 获取文件的绝对路径
        filePath = filedialog.askopenfilename(defaultextension=".bmp", filetypes=[('image', '.bmp'), ("BMP", ".bmp")])
        # 获取一下水印图片的宽和高，也就是解密的密钥
        img2 = Image.open(filePath, 'r')
        image_width = img2.size[0]
        image_height = img2.size[1]
        image_length = image_width * image_height * 24
        filePath2.set(filePath)

    def deEncry(self, img, length):
        # 获取图片的宽和高
        width = img.size[0]
        height = img.size[1]

        # 计数器
        count = 0
        # 结果文本，从图片中提取到的附加值（加密时附加在每个RGB通道后的二进制数值）
        wt = ""

        # 遍历图片
        for i in range(width):
            for j in range(height):
                # 获取像素点的值
                rgb = img.getpixel((i, j))

                # 提取R通道的附加值
                if count % 3 == 0:
                    count += 1
                    wt = wt + str(rgb[0] % 2)
                    if count == length:
                        break

                # 提取G通道的附加值
                if count % 3 == 1:
                    count += 1
                    wt = wt + str(rgb[1] % 2)
                    if count == length:
                        break

                # 提取B通道的附加值
                if count % 3 == 2:
                    count += 1
                    wt = wt + str(rgb[2] % 2)
                    if count == length:
                        break
            if count == length:
                break
        return wt

    def showImage(self, wt, img_width, img_height):
        str1 = []
        for i in range(0, len(wt), 8):
            # 以每8位为一组二进制，转换为十进制
            str1.append(int(wt[i:i + 8], 2))
        # 图片大于水印图片1个像素，便于对比
        img = Image.new("RGB", (img_width, img_height))
        flag = 0
        for m in range(0, img_width):
            for n in range(0, img_height):
                if flag == len(str1):
                    break
                img.putpixel((m, n), (str1[flag], str1[flag + 1], str1[flag + 2]))
                flag += 3
            if flag == len(str1):
                break

        # 要保存的文件名
        self.path = file_dir + '/' + name + '解密图像' + '.bmp'
        # 保存图片
        img.save(self.path)
        # 显示图片
        img.show()

    def ok(self):
        rgb_im1 = img1.convert('RGB')
        wt = self.deEncry(rgb_im1, image_length)
        self.showImage(wt, image_width, image_height)
        self.userinfo = [self.path]  # 设置数据
        self.destroy()  # 销毁窗口

    def cancel(self):
        self.userinfo = None  # 空！
        self.destroy()

# 主窗
class MyApp(tk.Tk):
  def __init__(self):
    super().__init__()
    self.title('数字水印：图像加密与检测')
    # 程序界面
    self.setupUI()

  def setupUI(self):
    # 生成画布
    frm = tk.Frame(self)
    frm.pack(padx=8, pady=8, ipadx=4, ipady=0)
    # 图像加密布局
    tk.Label(frm, text='加密图像文件路径：').grid(row=0, column=0, padx=5, pady=5)
    self.out_1 = StringVar()
    out1 = tk.Entry(frm, textvariable=self.out_1)
    out1.grid(row=0, column=1, padx=5, pady=5)
    tk.Button(frm, text="图像加密", command=self.setup_encry_config).grid(row=0, column=2, padx=5, pady=5)
    # 图像解密布局
    tk.Label(frm, text='解密图片文件路径：').grid(row=1, column=0, padx=5, pady=5)
    self.out_2 = StringVar()
    out2 = tk.Entry(frm, textvariable=self.out_2)
    out2.grid(row=1, column=1, padx=5, pady=5)
    tk.Button(frm, text="水印检测", command=self.setup_deEncry_config).grid(row=1, column=2,padx=5, pady=5)

  # 加密设置参数
  def setup_encry_config(self):
    # 接收弹窗的数据
    res = self.ask_userinfo_encry()
    if res is None: return
    # 更改参数
    self.path = res
    # 更新界面
    self.out_1.set(self.path)

  # 解密设置参数
  def setup_deEncry_config(self):
    # 接收弹窗的数据
    res = self.ask_userinfo_deEncry()
    if res is None: return
    # 更改参数
    self.path = res
    # 更新界面
    self.out_2.set(self.path)

  # 弹窗
  def ask_userinfo_encry(self):
    encry= MyDialogEncry()
    self.wait_window(encry)
    return encry.userinfo

  # 弹窗
  def ask_userinfo_deEncry(self):
    deEncry= MyDialogDeEncry()
    self.wait_window(deEncry)
    return deEncry.userinfo

if __name__ == '__main__':
  app = MyApp()
  app.mainloop()

